// 

// The central controller is responsible to manage the new
// tasks from the task generater, alloc gpus, routes, etc for
// the tasks, and manage the critical resources.

#ifndef SIMPLE_MODULES_CENTRAL_CONTROLLER_H_
#define SIMPLE_MODULES_CENTRAL_CONTROLLER_H_

#include <vector>
#include <string>
#include <map>
#include <unordered_map>
#include <utility>
#include "ib_cluster/base/ai_sim.h"
#include "ib_cluster/network_module/control/TopoManager.h"
#include "ib_cluster/task/allreduce.h"
#include "ib_cluster/route/route.h"

namespace ai_sim {

class CentralController : public omnetpp::cSimpleModule{
 private:
    int numMachine;
    int numSwitch;
    int numGpusperMachine;
    int numOxc;
    ClusterType clustertype;
    std::unordered_map<int, taskboard> taskboards;
    TopoManager *topos;
    AllreduceAlgorithm *allreduce_algo;
    RouteAlgorithm *route_algo;

 public:
    // The following redefined virtual function holds the algorithm.
    void initialize() override;
    void handleMessage(omnetpp::cMessage *msg) override;
    // Can be used in optical cluster or traditional cluster

    void alloc_gpu_resource(taskboard&);

    void gen_comm_pair(taskboard&, AllreduceAlgorithm*);

    void gen_flows(taskboard&, AllreduceAlgorithm*);

    void gen_route(taskboard&, RouteAlgorithm*);
    // Alloc flows of all stages according to taskboards
    void allocAllFlows(taskboard&);
    // Used when we need to alloc a batch of routes
    // to multi switches
    void allocRoutes(const map<int, RouteVec>&);

    // Alloc tasks from alloc gpu tp alloc flow and routes
    void alloc_task(const Task&);

    // Release the resources that a task used after the task
    // finished
    void try_release_task(const Task&, TaskMsg*);

    void get_statistics(const taskboard&);

    void finish() override;
};
}  // namespace ai_sim

#endif  // SIMPLE_MODULES_CENTRAL_CONTROLLER_H_
